import os
import pydicom
import pandas as pd
import numpy as np
from tqdm import tqdm

import torch
import torchvision

import torch.nn as nn
import torch.nn.functional as F
from torchvision.transforms.v2 import Compose, Resize, Normalize, ToTensor


class CNNModel(nn.Module):
    def __init__(self, mode="train", depth=16):
        super().__init__()
        if mode == "train":
            self.model = torchvision.models.resnet18(True)
        else:
            self.model = torchvision.models.resnet18(False)
        self.model.conv1 = nn.Conv2d(depth, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)  # todo 修改通道数
        self.model.fc = torch.nn.Linear(512, 256)
        self.bn1 = nn.BatchNorm1d(256)
        self.lu = nn.GELU()
        self.l1 = nn.Linear(256, 64)  # 使用清洗的数据
        self.bn2 = nn.BatchNorm1d(64)
        self.fc=nn.Linear(64,14)
        

    def forward(self, x):

        x = self.model(x)
        x = self.bn1(x)
        x = self.lu(x)
        x = self.l1(x)
        x = self.bn2(x)
        x=self.fc(x)
        return x
# 获得病人的路径 如果有两个路径需要进行均分结果

def fetch_img_paths(train_img_path):
    img_paths = []
    print('Scanning directories...')
    for patient in tqdm(os.listdir(train_img_path)):
        for scan in os.listdir(os.path.join(train_img_path, patient)):
            scan_path = os.listdir(os.path.join(train_img_path, patient, scan))
            scans = [
                [filename[:-4],
                    os.path.join(train_img_path, patient, scan, filename)]
                for filename in scan_path
                if filename.endswith('.dcm') and 10 <= int(filename[:-4]) <= 1500
            ]
            scans = sorted(scans, key=lambda x: int(x[0]))
            img_path = np.array(scans)
            img_paths.append(img_path[:, 1].tolist())
            min_length = min(min_length, len(scans))
    return img_paths


def select_images_p(numbers):
    samples = np.random.normal(loc=numbers/2, scale=1, size=numbers)
    sample_sum = np.sum(samples)
    normalized_samples = samples / sample_sum
    return normalized_samples.tolist()


def select_elements_with_spacing(input_list,depth):  # 读取40张图片
    length = len(input_list)
    if length <= depth:
        return np.random.choice(input_list,depth)
    else:
        indexs = np.linspace(0, length - 1, depth).astype(int)
    return [input_list[index] for index in indexs]


def preprocess_jpeg(dicm_img_path):
    dicoms = clip_rescale_extract_dicom_image(pydicom.read_file(dicm_img_path))
    # dicoms = pydicom.dcmread(dicm_img_path).pixel_array
    return dicoms


def clip_rescale_extract_dicom_image(dicom_ds):
    image = dicom_ds.pixel_array
    # find rescale params
    if ("RescaleIntercept" in dicom_ds) and ("RescaleSlope" in dicom_ds):
        intercept = float(dicom_ds.RescaleIntercept)
        slope = float(dicom_ds.RescaleSlope)
    center = int(dicom_ds.WindowCenter)
    width = int(dicom_ds.WindowWidth)
    low = center - width / 2
    high = center + width / 2
    image = (image * slope) + intercept
    image = np.clip(image, low, high)
    image = (image / np.max(image) * 255).astype(np.int16)
    return image


if __name__ == "__main__":
    # 使用compose进行数据增强
    depth = 16
    transform = Compose([
        Resize((224, 224), antialias=True),  # 进行重新训练  先用小的分辨率训练再使用大的分辨率训练
        Normalize([0.5], [0.5]),
        ToTensor(),
    ])

    # 定义模型
    model = CNNModel(mode="test", depth=depth)

    state_dict = torch.load('checkpoints/16_128_a_0.1_0.27199_model_78.pth',map_location=torch.device('cpu')) 
    model.load_state_dict(state_dict=state_dict['state_dict'],strict=False)
    model = model.cuda()
    model.eval()

    # 进行测试
    test_img_path = 'train_test_data/datasets/test_images'
    # test_pids = pd.read_csv('train_test_data/datasets/sample_submission.csv')['patient_id'].values
    test_pids =os.listdir(test_img_path)

    pred_list = []
    for pid in test_pids[:]:
        print(pid)
        preds = torch.zeros(1, 13)
        counts = 0
        for scan in os.listdir(os.path.join(test_img_path, str(pid))):
            counts = counts+1
            scan_path = os.listdir(os.path.join(test_img_path, str(pid), scan))
            scans = [[filename[:-4], os.path.join(test_img_path, str(pid), scan, filename)]
                    for filename in scan_path
                    if filename.endswith('.dcm') and 10 <= int(filename[:-4]) <= 1500]
            scans = sorted(scans, key=lambda x: int(x[0]))
            img_path = np.array(scans)
            indexs = select_elements_with_spacing(img_path[:, 1],depth)
            patient_id = indexs[0].split('/')[-3]
            images = []

            for d in indexs:
                image = preprocess_jpeg(d)
                images.append(image)
            images = np.stack(images)/1.0
            imgs = torch.tensor(images, dtype = torch.float).unsqueeze(0)
            imgs = transform(imgs)
            with torch.no_grad():
                imgs = imgs.cuda()
                output = model(imgs)[:, :-1]
                pred = torch.sigmoid(output).data.cpu().numpy()
            preds += torch.tensor(pred)
        pred_list.append(preds/(counts+0.005))

    pred_df = pd.DataFrame(np.vstack(pred_list))
    pred_df.columns = ['bowel_healthy', 'bowel_injury', 'extravasation_healthy',
                    'extravasation_injury', 'kidney_healthy', 'kidney_low', 'kidney_high',
                    'liver_healthy', 'liver_low', 'liver_high', 'spleen_healthy',
                    'spleen_low', 'spleen_high']

    pred_df['patient_id'] = test_pids[:]

    pred_df[['patient_id', 'bowel_healthy', 'bowel_injury', 'extravasation_healthy',
            'extravasation_injury', 'kidney_healthy', 'kidney_low', 'kidney_high',
            'liver_healthy', 'liver_low', 'liver_high', 'spleen_healthy',
            'spleen_low', 'spleen_high']].to_csv('submission.csv', index=None,float_format='%.4f')
